import argparse
import os
import glob
import cv2
import numpy as np
import copy


def original(i, j, ksize, img):
    # 找到矩阵坐标
    x1 = y1 = -ksize // 2
    x2 = y2 = ksize + x1
    temp = np.zeros(ksize * ksize)
    count = 0
    # 处理图像
    for m in range(x1, x2):
        for n in range(y1, y2):
            if i + m < 0 or i + m > img.shape[0] - 1 or j + n < 0 or j + n > img.shape[1] - 1:
                temp[count] = img[i, j]
            else:
                temp[count] = img[i + m, j + n]
            count += 1
    return temp

# 自定义最大值滤波器最小值滤波器
def max_min_functin(ksize, img, flag):
    # flag: 如果是0就检测最大值，如果是1就检测最小值。
    img0 = copy.copy(img)
    for i in range(0, img.shape[0]):
        for j in range(2, img.shape[1]):
            temp = original(i, j, ksize, img0)
            if flag == 0:   # 设置flag参数，如果是0就检测最大值，如果是1就检测最小值。
                img[i, j] = np.max(temp)
            elif flag == 1:
                img[i, j] = np.min(temp)
    return img


def parse_args():
    parser = argparse.ArgumentParser(description='Using template matching for license plate character recognition')
    # add argument as you like
    parser.add_argument('--img-dir', default='/work/dn/week3/work_dirs/test_clipped_resize_char_edge',
                        help='License plate image directory to be recognized')
    parser.add_argument('--tmp-dir', default='/work/dn/week3/work_dirs/Train_resized_edge')
    parser.add_argument('--method', default='sqdiff', help='the match method')
    parser.add_argument('--filter', default=None, help='the method of generate feature')

    args = parser.parse_args()
    return args


method_dict = {
    'sqdiff': cv2.TM_SQDIFF_NORMED,
    'ccorr': cv2.TM_CCORR_NORMED,
    'ccoeff': cv2.TM_CCOEFF_NORMED
}

if __name__ == '__main__':
    args = parse_args()
    assert args.method in method_dict
    # ========用三种方法进行模板匹配========
    method = method_dict[args.method]
    print(f'parsing image from {args.img_dir}...use the method: {args.method}')

    # 首先读取十张模板。
    templates = []
    for i in range(10):
        tem_dir = os.path.join(args.tmp_dir, 'sobel-'+str(i)+'.jpg')
        tem_img = cv2.imread(tem_dir, cv2.IMREAD_GRAYSCALE)   # h, w
        if args.filter is not None:
            if args.filter == 'blur':
                tem_img = cv2.blur(tem_img, (3, 3))    # 3x3均值滤波
            elif args.filter == 'median':
                tem_img = cv2.medianBlur(tem_img, 3)
            elif args.filter == 'max':
                tem_img = max_min_functin(3, tem_img, 0)
            else:
                tem_img = max_min_functin(3, tem_img, 1)
        templates.append(tem_img)
    h, w = templates[0].shape

    all = 0           # 总样本数
    right = 0         # 识别正确的数目

    for test_img_dir in glob.glob(os.path.join(args.img_dir, '*.jpg')):
        gt = int(test_img_dir.split('/')[-1].split('.')[0].split('-')[-1])
        test_img = cv2.resize(cv2.imread(test_img_dir, cv2.IMREAD_GRAYSCALE), (w, h), interpolation=cv2.INTER_NEAREST)
        if args.filter is not None:
            if args.filter == 'blur':
                test_img = cv2.blur(test_img, (3, 3))    # 3x3均值滤波
            elif args.filter == 'median':
                test_img = cv2.medianBlur(test_img, 3)
            elif args.filter == 'max':
                test_img = max_min_functin(3, test_img, 0)
            elif args.filter == 'min':
                test_img = max_min_functin(3, test_img, 1)

        res = []   # 记录对十个模板的匹配程度
        for tem in templates:
            result = cv2.matchTemplate(test_img, tem, method)
            min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(result)
            if method == cv2.TM_SQDIFF_NORMED:
                res.append(min_val)
            else:
                res.append(max_val)

        if method == cv2.TM_SQDIFF_NORMED:
            pred = np.argmin(np.array(res))
        else:
            pred = np.argmax(np.array(res))

        all += 1
        if pred == gt:
            right += 1

        img_name = test_img_dir.split('/')[-1]
        print(f'for img {img_name}, pred is {pred}, gt is {gt}.')
    print(f'total sample number is {all}, and the accuracy is equal to: {right/all}')